import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import re
from lifelines import KaplanMeierFitter, CoxPHFitter
from lifelines.exceptions import ConvergenceError
from sklearn.utils import resample

# 设置绘图风格 - 使用非交互式后端
plt.switch_backend('Agg')  # 使用非交互式后端避免显示问题
plt.style.use('default')

# 设置中文字体显示（如果需要显示中文）
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 加载数据
def load_and_preprocess_data(file_path):
    try:
        male_data = pd.read_excel(file_path, sheet_name='男胎检测数据')
    except Exception as e:
        print(f"数据加载失败: {e}")
        return None

    # 解析孕周数据
    def parse_gestational_week(week_str):
        if isinstance(week_str, str):
            match = re.match(r"(\d+)w\+(\d+)", week_str)
            if match:
                weeks = int(match.group(1))
                days = int(match.group(2))
                return weeks + days / 7
        return None

    male_data['孕周'] = male_data['检测孕周'].apply(parse_gestational_week)

    # 筛选有效数据
    male_data = male_data.dropna(subset=['Y染色体浓度', '孕妇BMI', '孕周'])
    return male_data


# 计算生存数据
def calculate_survival_data(group):
    group = group.sort_values('孕周')
    y_conc = group['Y染色体浓度'].values
    weeks = group['孕周'].values
    bmi = group['孕妇BMI'].iloc[0]  # BMI假设不变

    # 寻找首次达标的时间点
    success_idx = np.where(y_conc >= 0.04)[0]

    if len(success_idx) > 0:
        first_idx = success_idx[0]
        if first_idx == 0:
            # 第一次检测就达标
            T = weeks[0]
            event = 1
        else:
            # 区间删失情况
            lower = weeks[first_idx - 1]
            upper = weeks[first_idx]
            T = (lower, upper)
            event = 1
    else:
        # 从未达标，右删失
        T = weeks[-1]
        event = 0

    return pd.Series({
        'bmi': bmi,
        'T': T,
        'event': event,
        'first_week': weeks[0],
        'last_week': weeks[-1]
    })


"""执行生存分析"""
# 加载和预处理数据
male_data = load_and_preprocess_data("附件.xlsx")
if not male_data is None:
    # 计算每个孕妇的生存数据
    survival_df = male_data.groupby('孕妇代码').apply(calculate_survival_data, include_groups=False).reset_index()

    # BMI分组
    bins = [20, 28, 32, 36, 40, np.inf]  # 使用100作为上限而不是无穷大
    labels = ['20-28', '28-32', '32-36', '36-40', '40+']
    #survival_df['bmi_group'] = pd.cut(survival_df['bmi'], bins=bins, labels=labels, right=False)
    survival_df['bmi_group'] = pd.cut(survival_df['bmi'], bins=bins, labels=labels, include_lowest=True)
    # 为简化处理，将区间删失数据转换为中点值
    def approximate_T(row):
        if isinstance(row['T'], tuple):
            return np.mean(row['T'])
        return row['T']

    survival_df['T_approx'] = survival_df.apply(approximate_T, axis=1)

    # 绘制Kaplan-Meier生存曲线
    plt.figure(figsize=(10, 6))
    km_estimators = {}

    for group in labels:
        group_data = survival_df[survival_df['bmi_group'] == group]
        if len(group_data) > 0:
            kmf = KaplanMeierFitter()
            kmf.fit(
                durations=group_data['T_approx'],
                event_observed=group_data['event'],
                label=f'BMI {group}'
            )
            kmf.plot_survival_function(ci_show=False)
            km_estimators[group] = kmf

    plt.title('各BMI组Y染色体浓度达标生存曲线')
    plt.xlabel('孕周')
    plt.ylabel('未达标概率')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('question two 生存曲线.png', dpi=300)
    plt.close()

    # 使用Turnbull方法处理区间删失数据
    print("使用Turnbull方法处理区间删失数据:")
    for group in labels:
        group_data = survival_df[survival_df['bmi_group'] == group]
        if len(group_data) > 0:
            intervals = []
            for _, row in group_data.iterrows():
                if row['event'] == 1:
                    if isinstance(row['T'], tuple):
                        intervals.append(row['T'])
                    else:
                        intervals.append((row['T'], row['T']))
                else:
                    intervals.append((row['T'], np.inf))

            print(f"BMI {group} 组 - 找到 {len(intervals)} 个区间")

    # 计算最佳检测时点
    def find_optimal_time(kmf, threshold=0.9):
        """找到90%孕妇已达标的时间点"""
        survival_func = kmf.survival_function_
        cumulative_incidence = 1 - survival_func.iloc[:, 0]  # 获取第一列
        success_time = cumulative_incidence[cumulative_incidence >= threshold].index.min()
        return success_time if not pd.isna(success_time) else None

    optimal_times = {}
    for group, kmf in km_estimators.items():
        optimal_time = find_optimal_time(kmf)
        if optimal_time is not None:
            optimal_times[group] = optimal_time
            print(f"BMI {group} 组最佳检测时点: {optimal_time:.2f} 周")

    # Cox比例风险模型
    cox_data = survival_df[['T_approx', 'event', 'bmi']].copy()
    cox_data = cox_data.dropna()

    try:
        cph = CoxPHFitter()
        cph.fit(cox_data, duration_col='T_approx', event_col='event')
        print("\nCox比例风险模型结果:")
        print(cph.summary)
    except ConvergenceError:
        print("\nCox比例风险模型未能收敛，跳过此部分")

    # Bootstrap置信区间
    print("\n计算最佳检测时点的Bootstrap置信区间:")
    n_bootstrap = 100  # 减少Bootstrap次数以提高速度
    bootstrap_results = {group: [] for group in labels}

    for i in range(n_bootstrap):
        # 有放回抽样
        bootstrap_sample = resample(survival_df, random_state=i)

        for group in labels:
            group_data = bootstrap_sample[bootstrap_sample['bmi_group'] == group]
            if len(group_data) > 5:  # 确保有足够样本
                kmf_boot = KaplanMeierFitter()
                kmf_boot.fit(
                    durations=group_data['T_approx'],
                    event_observed=group_data['event']
                )
                optimal_time = find_optimal_time(kmf_boot)
                if optimal_time is not None:
                    bootstrap_results[group].append(optimal_time)

    # 计算置信区间
    for group in labels:
        if bootstrap_results[group]:
            ci_low = np.percentile(bootstrap_results[group], 2.5)
            ci_high = np.percentile(bootstrap_results[group], 97.5)
            print(f"BMI {group} 组最佳检测时点95%置信区间: [{ci_low:.2f}, {ci_high:.2f}]")
        else:
            print(f"BMI {group} 组 - 没有足够的Bootstrap样本计算置信区间")

    # SIMEX校正
    print("\n进行SIMEX校正...")
    sigma_e = 0.01  # 假设测量误差标准差
    lambda_values = [0, 0.5, 1, 1.5, 2]
    simex_coefficients = []

    for lam in lambda_values:
        coefs = []
        for _ in range(20):  # 减少模拟次数以提高速度
            cox_data_sim = cox_data.copy()
            # 添加测量误差
            error = np.random.normal(0, np.sqrt(lam) * sigma_e, len(cox_data_sim))
            cox_data_sim['T_approx'] += error

            try:
                cph_sim = CoxPHFitter()
                cph_sim.fit(cox_data_sim, duration_col='T_approx', event_col='event')
                coefs.append(cph_sim.params_['bmi'])
            except ConvergenceError:
                continue  # 跳过未能收敛的模拟

        if coefs:  # 确保有有效结果
            simex_coefficients.append(np.mean(coefs))

    # 二次外推
    def quadratic_extrapolation(x, y):
        coefficients = np.polyfit(x, y, 2)
        return np.polyval(coefficients, -1)

    if simex_coefficients:
        corrected_coef = quadratic_extrapolation(lambda_values[:len(simex_coefficients)], simex_coefficients)
        print(f"SIMEX校正后的BMI系数: {corrected_coef:.4f}")
    else:
        print("SIMEX校正失败，没有有效的结果")


